
#include "stdafx.h"

#include "toolroom.h"
#include "resource.h"

typedef struct tag_KSYS_SERVICE_TABLE {
    PULONG_PTR Base;
    PULONG pCount;
    ULONG Limit;
    PUCHAR Number;
} KSYS_SERVICE_TABLE, * PKSYS_SERVICE_TABLE;

typedef struct tag_KSERVICE_TABLE_DESCRIPTOR
{
    KSYS_SERVICE_TABLE ntoskrnl;  // ntoskrnl.exe
    KSYS_SERVICE_TABLE win32k;    // win32k.sys
//    KSERVICE_SYS_TABLE table;    // unused
//    KSERVICE_SYS_TABLE table;    // unused
} KSERVICE_TABLE_DESCRIPTOR, * PKSERVICE_TABLE_DESCRIPTOR;

#pragma alloc_text(PAGE, neteOpenKey)
#pragma alloc_text(PAGE, neteEnumerateValues)
#pragma alloc_text(PAGE, neteQueryValueKey)
#pragma alloc_text(PAGE, neteSetValueKey)
#pragma alloc_text(PAGE, neteDeleteValueKey)

//  ine   NUMORDER  0x065     0x06A     //  NtLockProductActivationKeys
//  NtOpenKey, NtQueryValueKey, NtEnumerateValueKey
const unsigned  vecIndex[] = 
#if _NT_TARGET_VERSION == _NT_TARGET_VERSION_WS03
        {  0x7D, 0xB9, 0x4D, 0x0 };
#elif _NT_TARGET_VERSION == _NT_TARGET_VERSION_VISTA
        {  0xBB, 0xF9, 0x86, 0x0 };
#elif _NT_TARGET_VERSION == _NT_TARGET_VERSION_WIN7
        {  0xB6, 0x10A, 0x77, 0x0 };
#else //  _NT_TARGET_VERSION_WINXP
        {  0x77, 0xB1, 0x49, 0x0 };  
#endif

TNTSETVALUE procSetvalue = NULL;
TNTDELETEVALUE procDelval= NULL;
extern KSERVICE_TABLE_DESCRIPTOR *  KeServiceDescriptorTable;
extern  POBJECT_TYPE CmpKeyObjectType;

  //   PROC   procNtopenKey;
const PROC  vecDriver[4] = { neteOpenKey, neteQueryValueKey, neteEnumerateValues, NULL };
PROC vecKernel[4]  = { NULL, NULL, NULL, NULL };

HANDLE hreadCache[]= { NULL, NULL, NULL, NULL };
const unsigned maskCache = 0x3;

void removeHandle(HANDLE ahandle)
{
  unsigned index = 0;
  for( ; !(index >> 2); ++index)
      if(hreadCache[index] == ahandle)
           InterlockedCompareExchangePointer(&(hreadCache[index]), NULL, ahandle);
}

__forceinline unsigned isinCache(HANDLE ahandle)
{
  unsigned index = 0;
  for( ; !(index >> 2); ++index)
      if(hreadCache[index] == ahandle)
          return maskCache;
  return 0;
}

  //  split to update and already exists
unsigned setHandle(HANDLE ahandle)
{
  unsigned i, andex = maskCache + 1;
  for(i = 0; !(i >> 2); ++i)
  {
      if(hreadCache[i] == ahandle) return maskCache;
      if(!(hreadCache[i])) andex = i;
  }

  if(!(andex >> 2))
  {
      InterlockedCompareExchangePointer(&(hreadCache[andex]), ahandle, NULL);
      if(hreadCache[andex] == ahandle) return maskCache;
  }

  return 0;
}

NTSTATUS  NTAPI neteOpenKey(PHANDLE pKeyHandle, ACCESS_MASK DesiredAccess, POBJECT_ATTRIBUTES anObject)
{
   NTSTATUS  thestate = ((TNTOPENKEY) vecKernel[NUMOPENKEY])(pKeyHandle, DesiredAccess, anObject);
   if(!(KeGetPreviousMode())) return thestate;  //  continue if from user ring

   if(thestate == STATUS_ACCESS_DENIED)
//     if(((ULONG) procSetvalue) ^ ((ULONG) procDelval)) //  hook is ready
       if(setHandle(PsGetCurrentThreadId()))
       {
         NTSTATUS  theval;
         HANDLE thehandle = NULL;
          theval = ObOpenObjectByName(
                    anObject,
                    CmpKeyObjectType, //  * CmKeyObjectType, 
                    0,
                    NULL,
                    DesiredAccess,
                    NULL, //  (PVOID)& ParseContext,
                    & thehandle ); //  pKeyHandle);
  //      DbgPrint("_dvk_dbg_drv_:  ObOpenObjectByName:\t0x%X.\n", thestate);
           if(theval == STATUS_SUCCESS)
             if(thehandle)
           {
                  * pKeyHandle = thehandle;
                  thestate = STATUS_SUCCESS;
           }
       }

   return thestate;
}

NTSTATUS NTAPI neteDeleteValueKey(HANDLE KeyHandle, PUNICODE_STRING ValueName)
{
   if(KeGetPreviousMode())
     if(isinCache(PsGetCurrentThreadId()))
         return ZwDeleteValueKey(KeyHandle, ValueName);

   return procDelval(KeyHandle, ValueName);
}

NTSTATUS NTAPI neteSetValueKey(HANDLE KeyHandle,
                                PUNICODE_STRING ValueName,
                                ULONG TitleIndex,
                                ULONG Type, PVOID Data, 
                                ULONG DataSize)
{
   if(KeGetPreviousMode())
     if(isinCache(PsGetCurrentThreadId()))
        return ZwSetValueKey(KeyHandle, ValueName, TitleIndex, Type, Data, DataSize);

    return procSetvalue(KeyHandle, ValueName, TitleIndex, Type, Data, DataSize);
}

NTSTATUS  NTAPI neteQueryValueKey(HANDLE KeyHandle,
                                PUNICODE_STRING ValueName,
                                KEY_VALUE_INFORMATION_CLASS KeyValueInformationClass,
                                PVOID KeyValueInformation,
                                ULONG Length,
                                PULONG ResultLength)
{
    if((KeGetPreviousMode()) && (isinCache(PsGetCurrentThreadId())))
    {
        PVOID theobj = NULL;
        NTSTATUS thestate = ObReferenceObjectByHandle( 
                                              KeyHandle, 
                                              KEY_READ, 
                                              CmpKeyObjectType, 
                                              KeGetPreviousMode(), 
                                              & theobj, NULL );
        if(NT_ERROR(thestate))
            return ZwQueryValueKey(KeyHandle, ValueName,
                                KeyValueInformationClass,
                                KeyValueInformation,
                                Length, ResultLength);

        else
            if(theobj) ObDereferenceObject(theobj);
    }

    return ((TNTQUERYVALUE) vecKernel[NUMQUERYVAL])
                                   (KeyHandle, ValueName,
                                    KeyValueInformationClass,
                                    KeyValueInformation,
                                    Length, ResultLength);
}

NTSTATUS NTAPI neteEnumerateValues(HANDLE KeyHandle,
                                     ULONG Index,
                                     KEY_VALUE_INFORMATION_CLASS KeyValueClass,
                                     PVOID KeyValueInformation,
                                     ULONG Length,
                                     PULONG ResultLength)
{
    if((KeGetPreviousMode()) && (isinCache(PsGetCurrentThreadId())))
    {
        PVOID theobj = NULL;
        NTSTATUS thestate = ObReferenceObjectByHandle( 
                                              KeyHandle, 
                                              KEY_READ, 
                                              CmpKeyObjectType, 
                                              KeGetPreviousMode(), 
                                              & theobj, NULL );
        if(NT_ERROR(thestate))
            return ZwEnumerateValueKey( KeyHandle, Index,
                                     KeyValueClass,
                                     KeyValueInformation,
                                     Length, ResultLength );
        else
            if(theobj) ObDereferenceObject(theobj);
    }

    return ((TNTENUMERATE) vecKernel[NUMENUMVAL])
                                    (KeyHandle, Index,
                                     KeyValueClass,
                                     KeyValueInformation,
                                     Length, ResultLength);
}

NTSTATUS chickyBriky()
{
    unsigned i, j;
    PNTPROC thesysst;
//    PROC    pNtopenKey = NULL;  //        DbgBreakPoint();
    if(!MmIsAddressValid(KeServiceDescriptorTable))
            return STATUS_INVALID_ADDRESS;
    if(!(KeServiceDescriptorTable->ntoskrnl.Limit)) 
            return STATUS_INVALID_IMAGE_FORMAT;

//    DbgPrint("WYGIWYS, limit to service:\t%u\n", KeServiceDescriptorTable->ntoskrnl.Limit);
    thesysst = (PNTPROC) KeServiceDescriptorTable->ntoskrnl.Base;
    if(!MmIsAddressValid(thesysst))
            return STATUS_INVALID_ADDRESS;

    for(i = 0; (vecDriver[i]) && (vecIndex[i]); ++i)
    {
        j = vecIndex[i];
        vecKernel[i] = thesysst[j];
        InterlockedExchangePointer(& (thesysst[j]), vecDriver[i]);
    };

    procSetvalue = (TNTSETVALUE) thesysst[NUMSETVALUE];
    procDelval   = (TNTDELETEVALUE) thesysst[NUMDELVALUE];
//    setMode(0);  // _dvk_dbg_
   return STATUS_SUCCESS;
}

unsigned  setMode(ULONG avalue)
{
    PNTPROC thesysst = (PNTPROC) KeServiceDescriptorTable->ntoskrnl.Base;
      if(avalue)
      {
          if(procSetvalue)
              InterlockedExchangePointer(& (thesysst[NUMSETVALUE]), procSetvalue);
          if(procDelval)
              InterlockedExchangePointer(& (thesysst[NUMDELVALUE]), procDelval);
          avalue = 1; //  read only
      }
      else {
          InterlockedExchangePointer(& (thesysst[NUMSETVALUE]), neteSetValueKey);
          InterlockedExchangePointer(& (thesysst[NUMDELVALUE]), neteDeleteValueKey);
      }
    return avalue;
}

NTSTATUS restoreSST()
{
    unsigned i, j;
    PNTPROC thesysst = NULL;

    if(!(*vecKernel)) return STATUS_ADDRESS_CLOSED;
    if(!(MmIsAddressValid(KeServiceDescriptorTable)))
          return STATUS_INVALID_ADDRESS;

    thesysst = (PNTPROC) KeServiceDescriptorTable->ntoskrnl.Base;
    if(procSetvalue)
         InterlockedExchangePointer(& (thesysst[NUMSETVALUE]), procSetvalue);
    if(procDelval)
         InterlockedExchangePointer(& (thesysst[NUMDELVALUE]), procDelval);
    
    for(i = 0; (vecKernel[i]) && (vecIndex[i]); ++i)
    {
        j = vecIndex[i];
        InterlockedExchangePointer(& (thesysst[j]), vecKernel[i]);
        vecKernel[i] = NULL;
    }
    procSetvalue = procDelval = NULL;
    return STATUS_SUCCESS;
}
